#include <unordered_map>
#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <algorithm>
#include <cassert>
#include <utility>
#include <cmath>
#include <cstring>
#include <random>
using namespace std;

inline FILE *fopen_check( const char *fname , const char *flag ){
    FILE *fp = fopen64( fname , flag );
    if( fp == NULL ){
        fprintf( stderr, "can not open file \"%s\"\n",fname );
        exit( -1 );
    }
    return fp;
}

class PairHashFunc {
public:
	size_t operator ()(const pair<int, int>& p) const {
		return (size_t) (p.first ^ p.second);
	}
};

bool greater_pred(const pair<double, int>& p1, const pair<double, int>& p2) {
    return p1.first > p2.first;
}

static const int num_items = 6095;
static const int num_users = 2320895;
static int AP_K = 3;

// user --> list of <pred, rating>

typedef unordered_map<int, vector<pair<double, int> > > UserRecMap;
typedef unordered_map<int, pair<double, int> > UserStatMap;

// vector<pair<double, int> > user_rec[num_users];
// vector<int> user_label[num_users];
// pair<double, int> stat[num_users];

UserRecMap user_rec;
unordered_map<int, vector<int> > user_label;
UserStatMap  stat;

mt19937 rand_gen(444);

void load_user_labels(const char* fname) {
	FILE* labelf = fopen_check(fname, "r");
	int uid, k, c;
	while(fscanf(labelf, "%d %d", &uid, &k) == 2) {
		for(int i=0; i<k; ++i) {
			assert( fscanf(labelf, "%d", &c) == 1 );
			user_label[uid].push_back(c);
		}
	}
	fclose(labelf);
}

int main(int argc, char* argv[]) {
    if (argc < 4) {
        printf("usage:<pred> <test> <wlist> [<user_label>] [K]\n");
		printf("\tnote that only <pred> is filtered!\n");
        return -1;
    }
    srand(444);

	FILE* predf = fopen_check(argv[1], "r");
	FILE* testf = fopen_check(argv[2], "r");
	FILE* wlistf = fopen_check(argv[3], "r");

	if (argc >= 5) {
		load_user_labels(argv[4]);
	}

	if (argc >= 6) {
		AP_K = atoi(argv[5]);
	}
	

	float pred;
	int uid, iid, rt, valid;
	long long t;
	while(fscanf(testf, "%d %d %d %lld", &uid, &iid, &rt, &t) == 4) {
		assert( fscanf(wlistf, "%d", &valid) == 1 );
		if (valid != 1) {
			continue;
		}

		assert( fscanf(predf, "%f", &pred) == 1 );		
		user_rec[uid].push_back(make_pair(pred, rt));

	}
	
	for(UserRecMap::iterator p = user_rec.begin();
		p != user_rec.end(); ++p) {
		// <pred, rt>
		vector<pair<double, int> >& rec = p->second;
		random_shuffle(rec.begin(), rec.end());
		partial_sort(rec.begin(), min(rec.end(), rec.begin() + AP_K), rec.end(), greater_pred);

		double sum_hit = 0.0;
		double sum_ap = 0.0;
		for(int i=0; i<AP_K && i < (int) rec.size(); ++i) {
			if (rec[i].second == 1) {
				sum_hit += 1;
				sum_ap += sum_hit / (i+1);
			}
		}

		int tot_hit = 0;
		for(vector<pair<double, int> >::iterator q = rec.begin();
			q != rec.end(); ++q) {
			tot_hit += q->second == 1 ? 1 : 0;
		}
		
		vector<int>& labels = user_label[p->first];
		if (labels.empty()) labels.push_back(0);

		for(vector<int>::iterator p = labels.begin();
			p != labels.end(); ++p) {
			int label = *p;

			pair<double, int>& ustat = stat[label];

			if (tot_hit > 0) {
				ustat.first += sum_ap / tot_hit;
				ustat.second += 1;
			}
		}
	}

	fclose(predf);
	fclose(testf);
	fclose(wlistf);
	
	printf("AP STAT @%d:\n", AP_K);


	vector<int> labels;
	for(UserStatMap::iterator p = stat.begin(); p != stat.end(); ++p) {
		labels.push_back(p->first);
	}
	sort(labels.begin(), labels.end());

	for(int i=0; i < (int) labels.size(); ++i) {
		int l = labels[i];
		pair<double, int>& st = stat[l];
		printf("\tlabel %d: tot = %d, average-AP = %.4f\n", l, st.second, st.first / st.second);
	}
}

	
